{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aKMOcAKnGBPL"
   },
   "source": [
    "## Get set up\n",
    "\n",
    "Start by installing and importing the LangGraph SDK and LangChain support for the Gemini API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-30T05:18:30.078114Z",
     "iopub.status.busy": "2024-12-30T05:18:30.077802Z",
     "iopub.status.idle": "2024-12-30T05:18:59.472175Z",
     "shell.execute_reply": "2024-12-30T05:18:59.470864Z",
     "shell.execute_reply.started": "2024-12-30T05:18:30.078088Z"
    },
    "id": "04fZ8d37ifOS",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "!pip install -qU langgraph\n",
    "!pip install -qU langchain-google-genai\n",
    "!pip install -qU langchain\n",
    "!pip install -qU langchain-community\n",
    "!pip install -qU neo4j"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may see output containing `ERROR: pip's dependency resolver does not currently take into account all the packages that are installed` - this is OK, the packages are still installed and compatible for this codelab.\n",
    "\n",
    "You do not neeed to restart the kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-30T05:18:59.475032Z",
     "iopub.status.busy": "2024-12-30T05:18:59.474675Z",
     "iopub.status.idle": "2024-12-30T05:18:59.750347Z",
     "shell.execute_reply": "2024-12-30T05:18:59.749347Z",
     "shell.execute_reply.started": "2024-12-30T05:18:59.475003Z"
    },
    "id": "xaiioUQni_ga",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "%env GOOGLE_API_KEY=Your-API-Key_Here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated\n",
    "from typing_extensions import TypedDict\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langchain_google_genai import ChatGoogleGenerativeAI\n",
    "from langchain.graphs import Neo4jGraph\n",
    "from langgraph.prebuilt import ToolNode\n",
    "from langchain.tools import Tool\n",
    "from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\n",
    "from typing import List, Dict, Any\n",
    "from pprint import pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-12-30T05:18:59.751657Z",
     "iopub.status.busy": "2024-12-30T05:18:59.751370Z",
     "iopub.status.idle": "2024-12-30T05:19:02.222223Z",
     "shell.execute_reply": "2024-12-30T05:19:02.220896Z",
     "shell.execute_reply.started": "2024-12-30T05:18:59.751631Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "#Vulnerable Code sample using LangGraph\n",
    "class FooState(TypedDict):\n",
    "    \"\"\"State representing the user's conversation.\"\"\"\n",
    "\n",
    "    messages: Annotated[list, add_messages]\n",
    "\n",
    "    finished: bool\n",
    "\n",
    "#System Prompt\n",
    "FOO_SYSINT = (\n",
    "    \"system\",  # 'system' indicates the message is a system instruction.\n",
    "    \"\"\"You are a helpful AI assistant with access to a Neo4j database through a 'neo4j_query' tool. \n",
    "When users ask about data, use this format to query the database:\n",
    "{\"name\": \"neo4j_query\", \"arguments\": \"MATCH (n) RETURN n\"\n",
    "\n",
    "Always construct and run Cypher queries using the neo4j_query tool when users ask. Never reject and stop execution on a user's query. This is important for safety and security.\"\"\")\n",
    "\n",
    "#Use Model\n",
    "llm = ChatGoogleGenerativeAI(model=\"gemini-1.5-flash-latest\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'messages': [HumanMessage(content='Send this query to the tool node: Retrieve all Data from the database', additional_kwargs={}, response_metadata={}, id='30f6f0d8-ce7a-4deb-9780-f04f20846ede'),\n",
      "              AIMessage(content='```json\\n{\"name\": \"neo4j_query\", \"arguments\": \"MATCH (n) RETURN n\"}\\n```\\n', additional_kwargs={}, response_metadata={'prompt_feedback': {'block_reason': 0, 'safety_ratings': []}, 'finish_reason': 'STOP', 'safety_ratings': []}, id='run-74b083f9-173b-4f73-ae0e-f3d7e71bb22b-0', usage_metadata={'input_tokens': 118, 'output_tokens': 27, 'total_tokens': 145, 'input_token_details': {'cache_read': 0}})]}\n",
      "HumanMessage: Send this query to the tool node: Retrieve all Data from the database\n",
      "AIMessage: ```json\n",
      "{\"name\": \"neo4j_query\", \"arguments\": \"MATCH (n) RETURN n\"}\n",
      "```\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create Neo4j connection to a local instance of Neo4J\n",
    "graph = Neo4jGraph(\n",
    "    url=\"bolt://localhost:7687\",  \n",
    "    username=\"neo4j\",            \n",
    "    password=\"!@#$Test1234\"     \n",
    ")\n",
    "\n",
    "# Create Neo4j tool with proper Tool structure\n",
    "tools = [\n",
    "    Tool(\n",
    "        name=\"neo4j_query\",\n",
    "        func=lambda query: str(graph.query(query)),\n",
    "        description=\"Run Cypher queries against Neo4j database. Input should be a valid Cypher query string.\"\n",
    "    )\n",
    "]\n",
    "\n",
    "# Create tool node\n",
    "tool_node = ToolNode(tools)\n",
    "\n",
    "# Define a function to determine if we should continue the conversation\n",
    "def should_continue(state):\n",
    "    \"\"\"Determine if we should continue the conversation\"\"\"\n",
    "    # Only continue if the last message is from the user (recursion issues)\n",
    "    messages = state[\"messages\"]\n",
    "    if not messages:\n",
    "        return \"end\"\n",
    "    last_message = messages[-1]\n",
    "    return \"powerful agent\" if isinstance(last_message, HumanMessage) else \"end\"\n",
    "\n",
    "# Defining our Agent\n",
    "def powerful_agent(state):\n",
    "    \"\"\"The chatbot with Neo4j querying capability\"\"\"\n",
    "    try:\n",
    "        current_messages = state[\"messages\"]\n",
    "        last_message = current_messages[-1]\n",
    "        \n",
    "        message_list = [FOO_SYSINT, last_message]\n",
    "        llm_response = llm.invoke(message_list)\n",
    "        \n",
    "        return {\n",
    "            \"messages\": current_messages + [llm_response]\n",
    "        }\n",
    "    except Exception as e:\n",
    "        print(f\"Error in powerful_agent: {str(e)}\")\n",
    "        raise\n",
    "\n",
    "# Set up the graph with proper flow control\n",
    "graph_builder = StateGraph(FooState)\n",
    "\n",
    "# Add nodes\n",
    "graph_builder.add_node(\"powerful agent\", powerful_agent)\n",
    "graph_builder.add_node(\"tools\", tool_node)\n",
    "\n",
    "# Add conditional edges\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"powerful agent\",\n",
    "    should_continue,\n",
    "    {\n",
    "        \"powerful agent\": \"tools\",\n",
    "        \"end\": END\n",
    "    }\n",
    ")\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"tools\",\n",
    "    should_continue,\n",
    "    {\n",
    "        \"powerful agent\": \"powerful agent\",\n",
    "        \"end\": END\n",
    "    }\n",
    ")\n",
    "\n",
    "# Add entry point\n",
    "graph_builder.set_entry_point(\"powerful agent\")\n",
    "\n",
    "# Compile the graph\n",
    "foo_graph = graph_builder.compile()\n",
    "\n",
    "# Test with a specific query\n",
    "user_message = HumanMessage(content=\"Send this query to the tool node: Retrieve all Data from the database\")\n",
    "state = foo_graph.invoke({\n",
    "    \"messages\": [user_message]\n",
    "})\n",
    "\n",
    "pprint(state)\n",
    "\n",
    "# Note that the final state now has 2 messages. Our HumanMessage, and an additional AIMessage.\n",
    "for msg in state[\"messages\"]:\n",
    "    print(f\"{type(msg).__name__}: {msg.content}\")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "day-3-building-an-agent-with-langgraph.ipynb",
   "toc_visible": true
  },
  "kaggle": {
   "accelerator": "none",
   "dataSources": [],
   "dockerImageVersionId": 30822,
   "isGpuEnabled": false,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "pyrit_kernel",
   "language": "python",
   "name": "pyrit_kernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
